package com.agent.rag.client;

import com.agent.commons.constants.MilvusConstants;
import com.agent.commons.utils.JacksonUtil;
import com.agent.feign.entity.SearchEntity;
import com.agent.rag.entity.MilvusEntity;
import com.agent.rag.mapper.RagPictureMapper;
import com.google.gson.Gson;
import com.google.gson.JsonObject;
import io.milvus.param.LogLevel;
import io.milvus.v2.client.MilvusClientV2;
import io.milvus.v2.common.IndexParam;
import io.milvus.v2.service.vector.request.InsertReq;
import io.milvus.v2.service.vector.request.QueryReq;
import io.milvus.v2.service.vector.request.SearchReq;
import io.milvus.v2.service.vector.request.data.EmbeddedText;
import io.milvus.v2.service.vector.request.data.FloatVec;
import io.milvus.v2.service.vector.response.QueryResp;
import io.milvus.v2.service.vector.response.SearchResp;
import jakarta.annotation.Resource;
import lombok.extern.slf4j.Slf4j;
import org.springframework.ai.embedding.EmbeddingModel;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Component;
import org.springframework.stereotype.Service;

import java.util.*;

/**
 * Milvus操作
 */
@Component
@Slf4j
public class MilvusClient {
    private final MilvusClientV2 milvusClientV2;

    @Resource
    private RagPictureMapper ragPictureMapper;
    @Autowired
    private EmbeddingModel embeddingModel;


    /**
     * 稠密向量权重
     */
    private float dense=0.7f;
    /**
     * 稀疏向量权重
     */
    private float sparse=0.1f;
    /**
     * 构造方法
     */
    @Autowired
    public MilvusClient(MilvusClientV2 milvusClientV2) throws InterruptedException {
        this.milvusClientV2=milvusClientV2;
        milvusClientV2.useDatabase(MilvusConstants.DATABASE_NAME);
    }


    /**
     * 设置稠密向量权重
     * @param dense
     * @return
     */
    public MilvusClient setDense(float  dense){
        this.dense=dense;
        return this;
    }

    /**
     * 设置稀疏向量权重
     * @param sparse
     * @return
     */
    public MilvusClient setSparse(float  sparse){
        this.sparse=sparse;
        return this;
    }

    /**
     * 向Milvus中插入数据
     */
    public void insert(List<MilvusEntity> milvusEntity) {
        //构建Gson对象
        Gson gson=new Gson();
        //构建JsonObject对象
        List<JsonObject> insertList = milvusEntity.stream().map(e -> {
            return gson.fromJson(JacksonUtil.writeValueAsString(e), JsonObject.class);
        }).toList();
        //Milvus插入语句
        milvusClientV2.insert(InsertReq.builder()
                .collectionName(MilvusConstants.TEXT_COLLECTION)
                .data(insertList)
                .build());
    }

    /**
     * 混合向量搜索
     */
    public List<SearchEntity> get( String text, String userId) {
        log.error("text:{},userId:{}",text,userId);
        float[] denseVector = embeddingModel.embed(text);
        Map<Long, Float> hybridMap=new HashMap<>();
        List<Long> ids=new ArrayList<>();
        //混合搜索
        //1.获取密集向量搜索到的结果
        Map<Long, Float> searchDense = this.searchDense(denseVector, userId);
        //2.将对应结果先一步添加到对应混合搜索结果中
        for(Long id:searchDense.keySet()){
            //2.1构造稀疏向量搜索结果
            ids.add(id);
            //2.2构造混合向量集合
            hybridMap.put(id,searchDense.get(id)*dense);
        }
        //3.获取稀疏向量搜索到的结果
        Map<Long, Float> searchSparse = this.searchSparse(ids,text, userId);
        //4.将对应结果添加到对应混合搜索结果中
        for(Long id:searchSparse.keySet()){
            //根据权重进行计算
            float hybridScore = hybridMap.get(id)+ searchSparse.get(id) * sparse;
            //添加到混合搜索结果中
            hybridMap.put(id,hybridScore);
        }
        //5.混合搜索结果排序
        //5.1根据id进行排序
        ids.sort( new Comparator<Long>() {
            @Override
            public int compare(Long o1, Long o2) {
                return Float.compare(hybridMap.get(o2),hybridMap.get(o1));
            }
        });
        //6.查询数据库,获取对应路径范围
        return this.getByIds(ids,userId);
    }

    /**
     * 稀疏向量搜索
     */
    private Map<Long,Float> searchSparse(List<Long> ids,String text,String userId){
        Map<Long,Float> idMap=new HashMap<>();//构建携带相似度分数及数据id的Map集合
        //BM25文本向量
        EmbeddedText embeddedText=new EmbeddedText(text);
        //构建稀疏向量查询请求
        SearchReq req=SearchReq.builder()
                .collectionName(MilvusConstants.TEXT_COLLECTION)
                .data(Collections.singletonList(embeddedText))
                .metricType(IndexParam.MetricType.BM25) //BM25相似度计算
                .filter("userId==\""+userId+"\"&& id in"+ids) //从密集搜索结果中获取id信息，在其中查询
                .annsField("sparse_vector")//指定查询字段
                .outputFields(List.of("id"))
                .topK(20)
                .build();
        SearchResp search = milvusClientV2.search(req);
        List<List<SearchResp.SearchResult>> searchResults = search.getSearchResults();
        for (List<SearchResp.SearchResult> searchResult : searchResults) {
            for (SearchResp.SearchResult result : searchResult) {
                idMap.put((Long)result.getEntity().get("id"),result.getScore());
            }
        }
        return idMap;
    }

    /**
     * 密集向量搜索
     * @param dense
     * @param userId
     * @return
     */
    private Map<Long,Float> searchDense(float[] dense, String userId){
        Map<Long,Float> ids=new HashMap<>();//构建携带相似度分数及数据id的Map集合
        // 构建查询向量
        FloatVec queryVector = new FloatVec(dense);
        //构建Milvus查询请求
        SearchReq req = SearchReq.builder()
                .collectionName(MilvusConstants.TEXT_COLLECTION)
                .data(Collections.singletonList(queryVector))
                .metricType(IndexParam.MetricType.IP) //内积相似度计算
                .filter("userId == \""+userId+"\"") //指定分区键
                .annsField("dense_vector") //指定查询字段
                .outputFields(List.of("id"))
                .topK(20)
                .build();
        SearchResp search = milvusClientV2.search(req);
        for(List<SearchResp.SearchResult> searchResult : search.getSearchResults()){
            for(SearchResp.SearchResult result : searchResult){
                ids.put((Long)result.getEntity().get("id"),result.getScore());
            }
        }
        return ids;
    }

    /**
     * 根据id范围查询文本及图片数据
     */
    private List<SearchEntity> getByIds(List<Long> ids, String userId){
        List<MilvusEntity> entities = new ArrayList<>();
        //构建查询请求
        QueryReq req = QueryReq.builder()
                .collectionName(MilvusConstants.TEXT_COLLECTION)
                .filter("userId == \""+userId+"\"&& id in"+ids)
                .outputFields(List.of("text","photo"))
                .build();
        QueryResp queryResp = milvusClientV2.query(req);
        List<SearchEntity> list = queryResp.getQueryResults().stream().map(e -> {
            SearchEntity searchEntity = new SearchEntity();
            String text = (String) e.getEntity().get("text");
            searchEntity.setText(text);
            List<Object> photo = (List<Object>) e.getEntity().get("photo");
            searchEntity.setPicture(photo.stream().map(picture->{
                return ragPictureMapper.getPictureUrl(picture.toString());
            }).toList());
            return searchEntity;
        }).toList();
        return list;
    }
}
